from flask import Flask, request, jsonify
from flask_cors import CORS
import torch
import os
import json
from model import load_model
import logging

app = Flask(__name__)
CORS(app)

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 全局变量存储模型
model = None
tokenizer = None
device = None

def initialize_model(model_path='models/trained'):
    """初始化模型"""
    global model, tokenizer, device
    
    try:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logger.info(f"使用设备: {device}")
        
        if os.path.exists(model_path):
            model, tokenizer = load_model(model_path, device)
            logger.info("模型加载成功")
            return True
        else:
            logger.warning(f"模型路径不存在: {model_path}")
            return False
    except Exception as e:
        logger.error(f"模型加载失败: {str(e)}")
        return False

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查接口"""
    return jsonify({
        'status': 'healthy',
        'model_loaded': model is not None,
        'device': str(device) if device else None
    })

@app.route('/chat', methods=['POST'])
def chat():
    """对话接口"""
    global model, tokenizer
    
    if model is None or tokenizer is None:
        return jsonify({
            'error': '模型未加载',
            'code': 'MODEL_NOT_LOADED'
        }), 500
    
    try:
        data = request.get_json()
        
        if not data or 'message' not in data:
            return jsonify({
                'error': '缺少必要参数: message',
                'code': 'MISSING_PARAMETER'
            }), 400
        
        user_message = data['message'].strip()
        if not user_message:
            return jsonify({
                'error': '消息不能为空',
                'code': 'EMPTY_MESSAGE'
            }), 400
        
        # 生成参数
        max_length = data.get('max_length', 50)
        temperature = data.get('temperature', 0.8)
        top_k = data.get('top_k', 50)
        
        # 生成回复
        logger.info(f"用户输入: {user_message}")
        
        response = model.generate(
            tokenizer=tokenizer,
            prompt=user_message,
            max_length=max_length,
            temperature=temperature,
            top_k=top_k
        )
        
        logger.info(f"AI回复: {response}")
        
        return jsonify({
            'response': response,
            'input': user_message,
            'parameters': {
                'max_length': max_length,
                'temperature': temperature,
                'top_k': top_k
            }
        })
        
    except Exception as e:
        logger.error(f"对话生成失败: {str(e)}")
        return jsonify({
            'error': f'对话生成失败: {str(e)}',
            'code': 'GENERATION_ERROR'
        }), 500

@app.route('/model/info', methods=['GET'])
def model_info():
    """获取模型信息"""
    global model
    
    if model is None:
        return jsonify({
            'error': '模型未加载',
            'code': 'MODEL_NOT_LOADED'
        }), 500
    
    try:
        # 计算参数数量
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        
        return jsonify({
            'model_type': 'SimpleAIModel',
            'vocab_size': model.vocab_size,
            'embed_dim': model.embed_dim,
            'hidden_dim': model.hidden_dim,
            'num_layers': len(model.transformer_blocks),
            'max_seq_len': model.max_seq_len,
            'total_parameters': total_params,
            'trainable_parameters': trainable_params,
            'device': str(next(model.parameters()).device)
        })
        
    except Exception as e:
        logger.error(f"获取模型信息失败: {str(e)}")
        return jsonify({
            'error': f'获取模型信息失败: {str(e)}',
            'code': 'INFO_ERROR'
        }), 500

@app.route('/model/reload', methods=['POST'])
def reload_model():
    """重新加载模型"""
    data = request.get_json()
    model_path = data.get('model_path', 'models/trained') if data else 'models/trained'
    
    logger.info(f"重新加载模型: {model_path}")
    
    if initialize_model(model_path):
        return jsonify({
            'message': '模型重新加载成功',
            'model_path': model_path
        })
    else:
        return jsonify({
            'error': '模型重新加载失败',
            'code': 'RELOAD_FAILED'
        }), 500

@app.route('/train/status', methods=['GET'])
def training_status():
    """获取训练状态（占位符，实际项目中可以实现训练任务管理）"""
    return jsonify({
        'status': 'not_implemented',
        'message': '训练状态监控功能待实现'
    })

@app.errorhandler(404)
def not_found(error):
    return jsonify({
        'error': '接口不存在',
        'code': 'NOT_FOUND'
    }), 404

@app.errorhandler(500)
def internal_error(error):
    return jsonify({
        'error': '服务器内部错误',
        'code': 'INTERNAL_ERROR'
    }), 500

if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='AI模型API服务器')
    parser.add_argument('--host', type=str, default='0.0.0.0', help='服务器主机地址')
    parser.add_argument('--port', type=int, default=5000, help='服务器端口')
    parser.add_argument('--model_path', type=str, default='models/trained', help='模型路径')
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    
    args = parser.parse_args()
    
    # 初始化模型
    logger.info("正在初始化AI模型API服务器...")
    model_loaded = initialize_model(args.model_path)
    
    if not model_loaded:
        logger.warning("模型未加载成功，API服务器将启动但对话功能不可用")
        logger.info("请先训练模型或检查模型路径")
    
    # 启动服务器
    logger.info(f"启动API服务器: http://{args.host}:{args.port}")
    app.run(host=args.host, port=args.port, debug=args.debug)